{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import timeit\n",
    "\n",
    "from Bio import PDB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading PDB structure '1TUP'...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tiago_antao/anaconda3/lib/python3.6/site-packages/Bio/PDB/StructureBuilder.py:90: PDBConstructionWarning: WARNING: Chain A is discontinuous at line 6146.\n",
      "  PDBConstructionWarning)\n",
      "/home/tiago_antao/anaconda3/lib/python3.6/site-packages/Bio/PDB/StructureBuilder.py:90: PDBConstructionWarning: WARNING: Chain B is discontinuous at line 6147.\n",
      "  PDBConstructionWarning)\n",
      "/home/tiago_antao/anaconda3/lib/python3.6/site-packages/Bio/PDB/StructureBuilder.py:90: PDBConstructionWarning: WARNING: Chain C is discontinuous at line 6148.\n",
      "  PDBConstructionWarning)\n",
      "/home/tiago_antao/anaconda3/lib/python3.6/site-packages/Bio/PDB/StructureBuilder.py:90: PDBConstructionWarning: WARNING: Chain E is discontinuous at line 6149.\n",
      "  PDBConstructionWarning)\n",
      "/home/tiago_antao/anaconda3/lib/python3.6/site-packages/Bio/PDB/StructureBuilder.py:90: PDBConstructionWarning: WARNING: Chain F is discontinuous at line 6171.\n",
      "  PDBConstructionWarning)\n",
      "/home/tiago_antao/anaconda3/lib/python3.6/site-packages/Bio/PDB/StructureBuilder.py:90: PDBConstructionWarning: WARNING: Chain A is discontinuous at line 6185.\n",
      "  PDBConstructionWarning)\n",
      "/home/tiago_antao/anaconda3/lib/python3.6/site-packages/Bio/PDB/StructureBuilder.py:90: PDBConstructionWarning: WARNING: Chain B is discontinuous at line 6383.\n",
      "  PDBConstructionWarning)\n",
      "/home/tiago_antao/anaconda3/lib/python3.6/site-packages/Bio/PDB/StructureBuilder.py:90: PDBConstructionWarning: WARNING: Chain C is discontinuous at line 6453.\n",
      "  PDBConstructionWarning)\n"
     ]
    }
   ],
   "source": [
    "repository = PDB.PDBList()\n",
    "parser = PDB.PDBParser()\n",
    "repository.retrieve_pdb_file('1TUP', file_format='pdb', pdir='.')  # XXX\n",
    "p53_1tup = parser.get_structure('P 53', 'pdb1tup.ent')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<Atom ZN> [58.108 23.242 57.424]\n",
      "<Atom ZN> [60.108 17.981 75.931]\n",
      "<Atom ZN> [33.653  0.403 74.115]\n"
     ]
    }
   ],
   "source": [
    "zns = []\n",
    "for atom in p53_1tup.get_atoms():\n",
    "    if atom.element == 'ZN':\n",
    "        #print(atom, dir(atom), atom.mass, atom.element, atom.coord[0])\n",
    "        zns.append(atom)\n",
    "for zn in zns:\n",
    "        print(zn, zn.coord)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Suggest a pymol viewing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Try this in numba?\n",
    "def get_closest_atoms(pdb_struct, ref_atom, distance):\n",
    "    atoms = {}\n",
    "    rx, ry, rz = ref_atom.coord\n",
    "    for atom in pdb_struct.get_atoms():\n",
    "        if atom == ref_atom:\n",
    "            continue\n",
    "        x, y, z = atom.coord\n",
    "        my_dist = math.sqrt((x - rx)**2 + (y - ry)**2 + (z - rz)**2) \n",
    "        if my_dist < distance:\n",
    "            atoms[atom] = my_dist\n",
    "    return atoms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[58.108 23.242 57.424]\n",
      "C 3.4080117696286854 [57.77  21.214 60.142]\n",
      "S 2.3262243799594877 [57.065 21.452 58.482]\n",
      "C 3.4566537492335123 [58.886 20.867 55.036]\n",
      "C 3.064120559761192 [58.047 22.038 54.607]\n",
      "N 1.9918273537290707 [57.755 23.073 55.471]\n",
      "C 2.9243719601324525 [56.993 23.943 54.813]\n",
      "C 3.857729198122736 [61.148 25.061 55.897]\n",
      "C 3.62725094648044 [61.61  24.087 57.001]\n",
      "S 2.2789209624943494 [60.317 23.318 57.979]\n",
      "C 3.087214470667822 [57.205 25.099 59.719]\n",
      "S 2.2253158446520818 [56.914 25.054 57.917]\n",
      "\n",
      "[60.108 17.981 75.931]\n",
      "C 3.41769274437124 [57.593 15.783 75.207]\n",
      "S 2.3254721582053093 [58.586 17.082 74.42 ]\n",
      "C 3.4672070967122894 [62.272 17.174 73.345]\n",
      "C 3.1139134725185587 [62.061 18.615 73.59 ]\n",
      "N 2.0564599972249455 [61.366 19.056 74.71 ]\n",
      "C 2.985233217423681 [61.332 20.382 74.647]\n",
      "C 3.805126390272999 [62.573 18.263 78.816]\n",
      "C 3.1803200512467478 [61.521 17.136 78.652]\n",
      "S 2.2070404885225816 [61.287 16.447 76.993]\n",
      "C 3.2038921042012745 [57.624 18.417 77.907]\n",
      "S 2.242320906916762 [58.978 19.402 77.247]\n",
      "\n",
      "[33.653  0.403 74.115]\n",
      "N 3.8418381161044053 [32.62  -3.267 73.642]\n",
      "C 3.4269003358801373 [31.435 -1.557 72.388]\n",
      "S 2.3788014279880345 [32.942 -0.607 72.082]\n",
      "C 3.1575681467123577 [36.183 -0.469 72.439]\n",
      "C 2.985114661190097 [36.225  0.98  72.714]\n",
      "N 2.0959130049879975 [35.24   1.603 73.456]\n",
      "C 3.2037898879982705 [35.569  2.894 73.492]\n",
      "C 3.940504780036419 [35.474  0.462 77.609]\n",
      "C 3.4076262123188337 [34.442 -0.639 77.262]\n",
      "S 2.3204200227887206 [34.453 -1.233 75.553]\n",
      "C 3.076855853166918 [30.82   1.082 75.105]\n",
      "S 2.1315044160944145 [32.391  1.939 74.884]\n"
     ]
    }
   ],
   "source": [
    "for zn in zns:\n",
    "    print()\n",
    "    print(zn.coord)\n",
    "    atoms = get_closest_atoms(p53_1tup, zn, 4)\n",
    "    for atom, distance in atoms.items():\n",
    "        print(atom.element, distance, atom.coord)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 [0, 0, 0]\n",
      "2 [1, 0, 0]\n",
      "4 [11, 11, 12]\n",
      "8 [109, 113, 106]\n",
      "16 [523, 721, 487]\n",
      "32 [2381, 3493, 2053]\n",
      "64 [5800, 5827, 5501]\n",
      "128 [5827, 5827, 5827]\n"
     ]
    }
   ],
   "source": [
    "for distance in [1, 2, 4, 8, 16, 32, 64, 128]:\n",
    "    my_atoms = []\n",
    "    for zn in zns:\n",
    "        atoms = get_closest_atoms(p53_1tup, zn, distance)\n",
    "        my_atoms.append(len(atoms))\n",
    "    print(distance, my_atoms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "86.03363559998343\n"
     ]
    }
   ],
   "source": [
    "nexecs = 10\n",
    "print(timeit.timeit('get_closest_atoms(p53_1tup, zns[0], 4.0)',\n",
    "                    'from __main__ import get_closest_atoms, p53_1tup, zns',\n",
    "                    number=nexecs) / nexecs * 1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_closest_alternative(pdb_struct, ref_atom, distance):\n",
    "    atoms = {}\n",
    "    rx, ry, rz = ref_atom.coord\n",
    "    for atom in pdb_struct.get_atoms():\n",
    "        if atom == ref_atom:\n",
    "            continue\n",
    "        x, y, z = atom.coord\n",
    "        if abs(x - rx) > distance or abs(y - ry) > distance or abs(z - rz) > distance:\n",
    "            continue\n",
    "        my_dist = math.sqrt((x - rx)**2 + (y - ry)**2 + (z - rz)**2) \n",
    "        if my_dist < distance:\n",
    "            atoms[atom] = my_dist\n",
    "    return atoms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "31.851707400005576\n"
     ]
    }
   ],
   "source": [
    "print(timeit.timeit('get_closest_alternative(p53_1tup, zns[0], 4.0)',\n",
    "                    'from __main__ import get_closest_alternative, p53_1tup, zns',\n",
    "                    number=nexecs) / nexecs * 1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Standard\n",
      "85.08649739999328\n",
      "86.50681579999855\n",
      "86.79630599999655\n",
      "96.95437099999253\n",
      "96.21982420001132\n",
      "Optimized\n",
      "30.253444099980698\n",
      "32.69531210000878\n",
      "52.965772600009586\n",
      "142.53310030001103\n",
      "141.26269519999823\n"
     ]
    }
   ],
   "source": [
    "print('Standard')\n",
    "for distance in [1, 4, 16, 64, 128]:\n",
    "    print(timeit.timeit('get_closest_atoms(p53_1tup, zns[0], distance)',\n",
    "                        'from __main__ import get_closest_atoms, p53_1tup, zns, distance',\n",
    "                        number=nexecs) / nexecs * 1000)\n",
    "print('Optimized')\n",
    "for distance in [1, 4, 16, 64, 128]:\n",
    "    print(timeit.timeit('get_closest_alternative(p53_1tup, zns[0], distance)',\n",
    "                        'from __main__ import get_closest_alternative, p53_1tup, zns, distance',\n",
    "                        number=nexecs) / nexecs * 1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "#for interesting distances"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
